import time
import os
import requests
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait
from contextlib import contextmanager

# 常量定义
DOMAINS_FILE = "domains.txt"
HOSTS_FILE = "hosts"
COOKIE_ACCEPT_BUTTON_XPATH = '//*[@id="cookie_box_accept_button"]'
INPUT_BOX_XPATH = '//*[@id="q"]'
SUBMIT_BUTTON_XPATH = '//*[@id="s"]'
RESULT_TABLE_XPATH = '//*[@id="results"]/tbody/tr/td[2]'
HEADERS = {
    "User-Agent": "Mozilla/5.0"
}
DNS_CHECKER_BASE_URL = "https://dnschecker.org"


class DNSCheckerClient:
    """DNS查询客户端类"""

    def __init__(self, headless=True):
        self.headless = headless
        self.driver_v4 = None
        self.driver_v6 = None

    def __enter__(self):
        """上下文管理器入口"""
        self.setup_drivers()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口，确保驱动关闭"""
        self.close_drivers()

    def setup_drivers(self):
        """配置并初始化Selenium WebDriver"""
        self.driver_v4 = self._create_driver()
        self.driver_v6 = self._create_driver()

    def _create_driver(self):
        """创建WebDriver实例"""
        opt = Options()
        if self.headless:
            opt.add_argument("--headless")  # 无头浏览器模式
            opt.add_argument("--disable-gpu")  # 禁用GPU加速
        return webdriver.Edge(options=opt)

    def close_drivers(self):
        """关闭所有WebDriver实例"""
        if self.driver_v4:
            self.driver_v4.quit()
        if self.driver_v6:
            self.driver_v6.quit()


def write_hosts_file(v4_content, v6_content, file_path):
    """将获取的IPv4和IPv6地址写入到hosts文件"""
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(f"## get host IP at {datetime.now().strftime('%Y-%m-%d %H:%M')} ##\n")
        f.write("127.0.0.1\tlocalhost\n::1\tlocalhost ip6-localhost ip6-loopback\n")
        f.write("ff02::1\tip6-allnodes\nff02::2\tip6-allrouters\n")
        f.write(v6_content + '\n\n' + v4_content)


def process_ip_list(data_list, domain, ip_type):
    """处理获取的IP地址列表并返回格式化字符串"""
    result_lines = []

    for data_element in data_list:
        data = data_element.text.strip()
        if not data or data == '-':
            print(f'{domain} 尚未获取到有效的 {ip_type} 地址，跳过')
            continue

        print(f"{domain} 获取到的 {ip_type} 地址是\n{data}")
        ip_addresses = data.split("\n")

        for ip in ip_addresses:
            if ip.strip():  # 确保不是空行
                result_lines.append(f"{ip.strip()}\t{domain}")

        time.sleep(1)  # 减少延时时间

    return '\n'.join(result_lines) if result_lines else ''


def fetch_dns_records(driver, domain, ip_type):
    """使用Selenium获取指定域名的DNS记录"""
    try:
        url = f'{DNS_CHECKER_BASE_URL}/country/cn/#{ip_type}/'
        driver.get(url)

        # 处理Cookie按钮，如果有的话
        try:
            cookie_button = WebDriverWait(driver, 5).until(
                EC.element_to_be_clickable((By.XPATH, COOKIE_ACCEPT_BUTTON_XPATH))
            )
            cookie_button.click()
            print("成功点击cookie接受按钮")
        except Exception:
            print("未找到cookie接受按钮，继续执行后续操作")

        # 输入域名并提交查询
        input_box = WebDriverWait(driver, 5).until(
            EC.presence_of_element_located((By.XPATH, INPUT_BOX_XPATH))
        )
        input_box.clear()
        input_box.send_keys(domain)
        time.sleep(0.5)

        submit_button = WebDriverWait(driver, 5).until(
            EC.element_to_be_clickable((By.XPATH, SUBMIT_BUTTON_XPATH))
        )
        submit_button.click()
        time.sleep(1)

        # 获取结果列表
        data_list = WebDriverWait(driver, 10).until(
            EC.presence_of_all_elements_located((By.XPATH, RESULT_TABLE_XPATH))
        )
        return data_list
    except Exception as e:
        print(f"获取 {domain} 的 {ip_type} 记录失败: {e}")
        return []


def check_dnschecker_access():
    """检查 DNSChecker 网站是否可访问"""
    try:
        response = requests.get(f"{DNS_CHECKER_BASE_URL}/", headers=HEADERS, timeout=10)
        if response.status_code != 200:
            print("https://dnschecker.org/ 不可达，5秒后自动退出程序")
            time.sleep(5)
            return False
        return True
    except requests.RequestException as e:
        print(f"网络请求失败: {e}")
        time.sleep(5)
        return False


def get_domain_mapping(domain):
    """获取域名映射，处理特殊情况"""
    domain_mappings = {
        'api.themoviedb.org': 'api.tmdb.org'
    }
    return domain_mappings.get(domain, domain)


def process_domains(domains_file_path, driver_v4, driver_v6):
    """处理所有域名并返回IPv4和IPv6结果"""
    v4_results = ['## 这里是ipv4地址']
    v6_results = ['## 这里是ipv6地址']

    try:
        with open(domains_file_path, "r", encoding="UTF-8") as f_domains:
            domains = [line.strip() for line in f_domains if line.strip()]
    except FileNotFoundError:
        print(f"错误：找不到文件 {domains_file_path}")
        return '', ''

    for domain in domains:
        print(f"正在查询 {domain}")
        query_domain = get_domain_mapping(domain)

        # 获取 IPv4 记录
        v4_data = fetch_dns_records(driver_v4, query_domain, "A")
        if v4_data:
            v4_result = process_ip_list(v4_data, domain, "A")
            if v4_result:
                v4_results.append(v4_result)

        # 获取 IPv6 记录
        v6_data = fetch_dns_records(driver_v6, query_domain, "AAAA")
        if v6_data:
            v6_result = process_ip_list(v6_data, domain, "AAAA")
            if v6_result:
                v6_results.append(v6_result)

    return '\n'.join(v4_results), '\n'.join(v6_results)


def main():
    if not check_dnschecker_access():
        return  # 如果 DNSChecker 不可访问则退出程序

    # 使用上下文管理器确保资源正确释放
    with DNSCheckerClient() as client:
        path = os.getcwd()
        domains_file_path = os.path.join(path, DOMAINS_FILE)
        hosts_file_path = os.path.join(path, HOSTS_FILE)

        v4, v6 = process_domains(domains_file_path, client.driver_v4, client.driver_v6)

        if v4 and v6:
            write_hosts_file(v4, v6, hosts_file_path)
            print(f"写入文件完成，文件位置{hosts_file_path}")
        else:
            print("未能获取到有效的DNS记录")


if __name__ == "__main__":
    main()
